#include "kernel_operator.h"

constexpr uint32_t BUFFER_NUM = 2;   
constexpr uint32_t ALIGN_BYTES = 32; 

template <typename T>
class KernelMaskedFill
{
public:
    __aicore__ inline KernelMaskedFill() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR mask, GM_ADDR valueT, GM_ADDR z, 
                                uint32_t inputBytes, uint32_t formerNum, 
                                uint32_t formerLength, uint32_t formerTileNum, 
                                uint32_t formerTileLength, uint32_t formerLastTileLength,
                                uint32_t tailLength, uint32_t tailTileNum,
                                uint32_t tailTileLength, uint32_t tailLastTileLength,
                                uint32_t xSize, uint32_t sShape[], uint32_t xNumShapes,
                                uint32_t maskSize, uint32_t maskShape[], uint32_t maskNumShapes)
    {
        valueGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(valueT), 1);
        this->value = valueGm.GetValue(0);

        this->alignNum_x = ALIGN_BYTES / inputBytes;

        this->size = xSize;      
        this->osize = maskSize;    
        this->numshapes = xNumShapes;
        this->onumshapes = maskNumShapes;
        for (uint32_t i = 0; i < this->numshapes; i++) {
            this->shape[i] = sShape[i];
        }
        for (uint32_t i = 0; i < this->onumshapes; i++) {
            this->oshape[i] = maskShape[i];
        }

        uint32_t dimOffsetX = 1;
        dimOffsetXList[this->numshapes - 1] = 1;
        for (int32_t i = this->numshapes - 2; i >= 0; i--) {
            dimOffsetX *= this->shape[i + 1];
            dimOffsetXList[i] = dimOffsetX;
        }
        uint32_t dimOffsetMask = 1;
        dimOffsetMaskList[this->onumshapes - 1] = 1;
        for (int32_t i = this->onumshapes - 2; i >= 0; i--) {
            dimOffsetMask *= this->oshape[i + 1];
            dimOffsetMaskList[i] = dimOffsetMask;
        }
        this->alignStart = this->numshapes - this->onumshapes;

        uint32_t blockIdx = AscendC::GetBlockIdx();
        uint32_t gmStartIdx = 0U;
        uint32_t blockTotalElemAligned = 0U;

        if (blockIdx < formerNum) {
            blockTotalElemAligned = formerLength;
            gmStartIdx = blockTotalElemAligned * blockIdx;
            this->tileNum = formerTileNum;
            this->tileLength_x = formerTileLength;
            this->lastTileLength_x = formerLastTileLength;
        } else {
            blockTotalElemAligned = tailLength;
            gmStartIdx = formerLength * formerNum + blockTotalElemAligned * (blockIdx - formerNum);
            this->tileNum = tailTileNum;
            this->tileLength_x = tailTileLength;
            this->lastTileLength_x = tailLastTileLength;
        }

        this->gmStartIdx = gmStartIdx;
        if (this->gmStartIdx >= xSize) {
            this->localValidElem = 0U;
            return;
        }
        this->localValidElem = (this->gmStartIdx + blockTotalElemAligned > xSize) ? 
                               (xSize - this->gmStartIdx) : blockTotalElemAligned;

        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x) + this->gmStartIdx, this->localValidElem);
        maskGm.SetGlobalBuffer(reinterpret_cast<__gm__ bool*>(mask), this->osize);
        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(z) + this->gmStartIdx, this->localValidElem);

        uint32_t xBufSize = this->tileLength_x * inputBytes;
        pipe.InitBuffer(inQueueX, BUFFER_NUM, xBufSize);
        pipe.InitBuffer(inQueueM, BUFFER_NUM, xBufSize); // mask缓冲区大小与x一致
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, xBufSize);
    }

    __aicore__ inline void Process()
    {
        if (this->localValidElem == 0U || this->tileNum == 0U) return;
        for (uint32_t i = 0U; i < this->tileNum; i++) {
            CopyIn(i);  
            Compute(i);  
            CopyOut(i); 
        }
    }

private:
    __aicore__ inline uint32_t MapXGlobalIdxToMaskGlobalIdx(uint32_t xGlobalIdx) {
        uint32_t xIndice[128] = {0};
        uint32_t temp = xGlobalIdx;
        for (int32_t i = this->numshapes - 1; i >= 0; i--) {
            xIndice[i] = temp % this->shape[i];
            temp /= this->shape[i];
        }

        uint32_t maskIndice[128] = {0};
        for (int32_t i = 0; i < this->onumshapes; i++) {
            int32_t xDimIdx = this->alignStart + i; 
            uint32_t maskDimSize = this->oshape[i];

            maskIndice[i] = (maskDimSize == 1) ? 0 : xIndice[xDimIdx];
        }

        uint32_t maskGlobalIdx = 0;
        for (int32_t i = 0; i < this->onumshapes; i++) {
            maskGlobalIdx += maskIndice[i] * dimOffsetMaskList[i];
        }
        return maskGlobalIdx;
    }

    __aicore__ inline uint32_t AlignUp(uint32_t value, uint32_t align)
    {
        return (value == 0U) ? 0U : ((value + align - 1U) & ~(align - 1U));
    }

    __aicore__ inline uint32_t GetCopyCount(uint32_t tileIdx)
    {
        uint32_t tileLen = (tileIdx == this->tileNum - 1U) ? this->lastTileLength_x : this->tileLength_x;
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        uint32_t maxValid = this->localValidElem - offsetInCore;

        return (maxValid <= 0U) ? 0U : ((tileLen <= maxValid) ? tileLen : AlignUp(maxValid, this->alignNum_x));
    }

    __aicore__ inline void CopyIn(uint32_t tileIdx)
    {
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        if (offsetInCore >= this->localValidElem) {
            this->xCopyCountCache[tileIdx % BUFFER_NUM] = 0U;
            return;
        }

        AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
        AscendC::LocalTensor<bool> maskLocal = inQueueM.AllocTensor<bool>();

        uint32_t xCopyCount = GetCopyCount(tileIdx);
        this->xCopyCountCache[tileIdx % BUFFER_NUM] = xCopyCount;

        if (xCopyCount > 0U) {
            AscendC::DataCopy(xLocal, xGm[offsetInCore], xCopyCount);
        }
        // 这里可以用Duplicate接口？
        if (xCopyCount > 0U) {
            for (uint32_t i = 0; i < xCopyCount; i++) {
                uint32_t xGlobalIdx = this->gmStartIdx + offsetInCore + i;
                uint32_t maskGlobalIdx = MapXGlobalIdxToMaskGlobalIdx(xGlobalIdx);
                maskLocal.SetValue(i, maskGm.GetValue(maskGlobalIdx));
            }
        }
        inQueueX.EnQue(xLocal);
        inQueueM.EnQue(maskLocal);
    }

    __aicore__ inline void Compute(uint32_t tileIdx)
    {
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        uint32_t bufIdx = tileIdx % BUFFER_NUM;
        if (offsetInCore >= this->localValidElem) return;

        AscendC::LocalTensor<T> xLocal = inQueueX.DeQue<T>();
        AscendC::LocalTensor<bool> maskLocal = inQueueM.DeQue<bool>();
        AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();

        uint32_t xCopyCount = this->xCopyCountCache[bufIdx];
        uint32_t maxValid = this->localValidElem - offsetInCore;
        uint32_t computeCount = AscendC::Std::min(xCopyCount, maxValid);
        for (int32_t i = computeCount - 1; i >= 0; i--)
        {
            // zLocal.SetValue(i, maskLocal.GetValue(i) ? this->value : xLocal.GetValue(i));
            // int16_t/uint16_t/half/int32_t/uint32_t/float/bfloat16_t 适用Duplicate接口
            if constexpr (AscendC::Std::is_same<T, bfloat16_t>::value || AscendC::Std::is_same<T, float>::value ||
                          AscendC::Std::is_same<T, uint32_t>::value || AscendC::Std::is_same<T, int32_t>::value ||
                          AscendC::Std::is_same<T, half>::value || AscendC::Std::is_same<T, uint16_t>::value ||
                          AscendC::Std::is_same<T, int16_t>::value) {
                AscendC::Duplicate<T>(zLocal, maskLocal.GetValue(i) ? this->value : xLocal.GetValue(i), i + 1);
            }
            else {
                zLocal.SetValue(i, maskLocal.GetValue(i) ? this->value : xLocal.GetValue(i));
            }
        }

        outQueueZ.EnQue(zLocal);
        inQueueX.FreeTensor(xLocal);
        inQueueM.FreeTensor(maskLocal);
    }

    __aicore__ inline void CopyOut(uint32_t tileIdx)
    {
        uint32_t offsetInCore = tileIdx * this->tileLength_x;
        uint32_t bufIdx = tileIdx % BUFFER_NUM;
        if (offsetInCore >= this->localValidElem) return;

        AscendC::LocalTensor<T> zLocal = outQueueZ.DeQue<T>();
        uint32_t copyCount = GetCopyCount(tileIdx);
        if (copyCount > 0U) {
            AscendC::DataCopy(zGm[offsetInCore], zLocal, copyCount);
        }

        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;                               
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX, inQueueM; 
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;       
    AscendC::GlobalTensor<T> xGm;   
    AscendC::GlobalTensor<bool> maskGm;                
    AscendC::GlobalTensor<T> valueGm;                               
    AscendC::GlobalTensor<T> zGm;                                         

    // -------------------------- 广播变量 ------------------------------
    uint32_t size;                    // x总元素数
    uint32_t shape[128];              // x原始shape
    uint32_t numshapes;               // x维度数
    uint32_t osize;                   // mask总元素数
    uint32_t oshape[128];             // mask原始shape
    uint32_t onumshapes;              // mask维度数
    uint32_t dimOffsetXList[128];     // x的维度偏移表
    uint32_t dimOffsetMaskList[128];  // mask的维度偏移表
    int32_t alignStart;               // 广播对齐起始维度

    // -------------------------- 分片变量 ------------------------------
    uint32_t tileNum;                 // 当前Core的Tile总数
    uint32_t alignNum_x;              // x的对齐元素数
    uint32_t tileLength_x;            // x的常规Tile长度（对齐后）
    uint32_t lastTileLength_x;        // x的尾Tile长度（对齐后）
    uint32_t gmStartIdx;              // 当前Core的GM起始索引
    uint32_t localValidElem;          // 当前Core的有效元素数

    // -------------------------- 缓存变量 ------------------------------
    uint32_t xCopyCountCache[BUFFER_NUM] = {0U};  // x拷贝数缓存
    T value;                                      // 填充
};

extern "C" __global__ __aicore__ void masked_fill_custom(GM_ADDR x, GM_ADDR mask, GM_ADDR value, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelMaskedFill<DTYPE_X> op;
    op.Init(x, mask, value, z,  
            tiling_data.inputBytes, tiling_data.formerNum, 
            tiling_data.formerLength, tiling_data.formerTileNum, 
            tiling_data.formerTileLength, tiling_data.formerLastTileLength,

            tiling_data.tailLength, tiling_data.tailTileNum, 
            tiling_data.tailTileLength, tiling_data.tailLastTileLength,

            tiling_data.xSize, tiling_data.sShape, tiling_data.xNumShapes,
            tiling_data.maskSize, tiling_data.maskShape, tiling_data.maskNumShapes);  
    op.Process();
}